package com.tezign.intelligence.mini.dam.aifamilyconsultant.controller;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.bean.ResponseBean;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.bean.resp.SessionVO;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.constant.Constant;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.entity.*;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.interceptor.FamilyStore;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.interceptor.UserStore;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.service.*;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.util.ChatGPTUtil;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.util.ChatUtil;
//import com.tezign.intelligence.mini.dam.aifamilyconsultant.util.DeepSeekUtil;
import com.tezign.intelligence.mini.dam.aifamilyconsultant.util.DeepSeekUtil;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import io.reactivex.Flowable;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;

import static io.lettuce.core.pubsub.PubSubOutput.Type.message;

/**
 * <p>
 * 单次会话信息 前端控制器
 * </p>
 *
 * @author six
 * @since 2024-11-16
 */
@Api(tags = "单次会话管理")
@Slf4j
@CrossOrigin("*")
@RestController
@RequiredArgsConstructor
@RequestMapping(Constant.APPLICATION_NAME + "/session")
public class SessionController {

    private final ISessionService sessionService;

    private final IUserSessionService userSessionService;

    private final IFamilyService familyService;

    private final ILargeModelService largeModelService;


    // 获取单次会话信息
    @GetMapping("/get")
    @ApiOperation(value = "获取单次会话信息")
    public ResponseBean<List<SessionVO>> get(@RequestParam("userSessionId") Long userSessionId) {
        return new ResponseBean<List<SessionVO>>().success(getListSession(userSessionId).stream().map(x -> BeanUtil.copyProperties(x, SessionVO.class)).collect(Collectors.toList()));
    }

    // 开始单次会话
    @PostMapping("/start/{userSessionId}")
    @ApiOperation(value = "开始单次会话")
    public SseEmitter start(
            @PathVariable("userSessionId") Long userSessionId,
            @RequestBody String content
    ) {
        if (StringUtils.isEmpty(content)) {
            throw new RuntimeException("内容不能为空");
        }
        List<Session> sessions = getListSession(userSessionId);
        List<Message> messages = sessions.stream().map(x -> BeanUtil.copyProperties(x, Message.class)).collect(Collectors.toList());
        messages.add(Message.builder().role(Role.USER.getValue()).content(content).build());
        Session se = new Session();
        se.setUserSessionId(userSessionId);
        se.setRole(Role.USER.getValue());
        se.setContent(content);
        sessionService.save(se);
        SseEmitter emitter = new SseEmitter();
        try {
            Flowable<GenerationResult> result = ChatUtil.streamCall(messages);
            result.subscribe(
                    message -> {
                        try {
                            String msg = message.getOutput().getChoices().get(0).getMessage().getContent();
                            String requestId = message.getRequestId();
                            Session one = sessionService.lambdaQuery().eq(Session::getUserSessionId, userSessionId).eq(Session::getRequestId, requestId).one();
                            if (one != null) {
                                one.setContent(one.getContent() + msg);
                                sessionService.updateById(one);
                            } else {
                                Session save = new Session();
                                String role = message.getOutput().getChoices().get(0).getMessage().getRole();
                                save.setUserSessionId(userSessionId);
                                save.setRole(role);
                                save.setRequestId(requestId);
                                save.setContent(msg);
                                sessionService.save(save);
                            }
                            emitter.send(msg);
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    },
                    emitter::completeWithError, // 如果出错，终止流
                    emitter::complete           // 完成整个流
            );
        } catch (Exception e) {
            emitter.completeWithError(e);
        }
        return emitter;
    }
//
//	@PostMapping("/question/{userSessionId}")
//	@ApiOperation(value = "开始回答问题")
//	public ResponseBean<String> question(
//			@PathVariable("userSessionId") Long userSessionId,
//			@RequestBody String content
//	) {
//		if (StringUtils.isEmpty(content)) {
//			throw new RuntimeException("内容不能为空");
//		}
//		ResponseBean<String> responseBean = new ResponseBean<>();
//		FamilyMember familyMember = familyMemberService.lambdaQuery().eq(FamilyMember::getFamilyId, FamilyStore.getFamilyId()).eq(FamilyMember::getUserId, UserStore.getUserId()).last("limit 1").one();
//		if (familyMember == null) {
//			throw new RuntimeException("用户不存在");
//		}
//		List<UserSessionQuestion> list = userSessionQuestionService.lambdaQuery().eq(UserSessionQuestion::getRole, familyMember.getRole()).orderByDesc(UserSessionQuestion::getSort).list();
//		if (CollectionUtils.isEmpty(list)) {
//			return responseBean.success("");
//		}
//		for (UserSessionQuestion userSessionQuestion : list) {
//			String key = String.format(QUESTION_KEY, userSessionId);
//			String value = redisTemplate.opsForValue().get(key);
//			if (StringUtils.isEmpty(value) || Integer.parseInt(value) > userSessionQuestion.getSort()) {
//				redisTemplate.opsForValue().set(key, userSessionQuestion.getSort().toString(), Duration.ofDays(30L));
//				String question = userSessionQuestion.getQuestion();
//				Session save = new Session();
//				save.setUserSessionId(userSessionId);
//				save.setRole(ASSISTANT.value());
//				save.setContent(question);
//				sessionService.save(save);
//				return responseBean.success(question);
//			}
//
//		}
//
//		return responseBean.success("");
//
//	}


    // 开始单次会话
    @PostMapping("/startDouBao/{userSessionId}")
    @ApiOperation(value = "开始单次会话")
    public SseEmitter startDouBao(
            @PathVariable("userSessionId") Long userSessionId,
            @RequestBody String content
    ) {
        if (StringUtils.isEmpty(content)) {
            throw new RuntimeException("内容不能为空");
        }

        List<Session> sessions = getListSession(userSessionId);
        List<ChatMessage> messages = sessions.stream().map(x -> {
            ChatMessage chatMessage = new ChatMessage();
            chatMessage.setRole(ChatMessageRole.valueOf(x.getRole().toUpperCase()));
            chatMessage.setContent(x.getContent());
            return chatMessage;
        }).collect(Collectors.toList());
        messages.add(ChatMessage.builder().role(ChatMessageRole.USER).content(content).build());
        Session se = new Session();
        se.setUserSessionId(userSessionId);
        se.setRole(Role.USER.getValue());
        se.setContent(content);
        sessionService.save(se);

        SseEmitter emitter = new SseEmitter();
        try {
            LargeModel dbLargeModel = largeModelService.getOne(new LambdaQueryWrapper<LargeModel>().eq(LargeModel::getEnable, "启用"));
            if (dbLargeModel != null) {
                chatGPTResult(emitter, messages, userSessionId, dbLargeModel.getModel());
//                switch (dbLargeModel.getCode()){
//                    case "01":
//                        // 豆包
//                        douBaoChatResult(emitter, messages, userSessionId);
//                        break;
//                    case "02":
//                        // deepseek
//                        deepSeekChatResult(emitter, messages, userSessionId);
//                        break;
//                    case "03":
//                        // chatGPT
//                        chatGPTResult(emitter, messages, userSessionId);
//                        break;
//                }
            }else {
                douBaoChatResult(emitter, messages, userSessionId);
            }
        } catch (Exception e) {
            emitter.completeWithError(e);
        }
        return emitter;
    }


    @PostMapping("/addSession/{userSessionId}")
    @ApiOperation(value = "补充预回答问题")
    public ResponseBean<Boolean> addSession(@PathVariable("userSessionId") Long userSessionId, @RequestBody List<SessionVO> sessionVOList) {
        if (CollUtil.isEmpty(sessionVOList)) {
            throw new RuntimeException("参数不能为空");
        }
        List<Session> res = sessionVOList.stream().map(x -> {
            Session session = new Session();
            session.setUserSessionId(userSessionId);
            session.setRole(x.getRole());
            session.setContent(x.getContent());
            return session;
        }).collect(Collectors.toList());
        sessionService.saveBatch(res);
        return new ResponseBean<Boolean>().success(true);
    }


    private List<Session> getListSession(Long userSessionId) {
        Long userId = UserStore.getUserId();
        Long familyId = FamilyStore.getFamilyId();
        log.info("获取单次会话信息, userId: {}, familyId: {}, userSessionId: {}", userId, familyId, userSessionId);
        UserSession one = userSessionService.lambdaQuery().eq(UserSession::getUserId, userId).eq(UserSession::getFamilyId, familyId).eq(UserSession::getId, userSessionId).one();

        Family family = familyService.getById(familyId);
        if (family == null) {
            throw new RuntimeException("家庭不存在");
        }

        if (one == null) {
            throw new RuntimeException("用户会话不存在");
        }
        Long id = one.getId();
        return sessionService.list(Wrappers.<Session>lambdaQuery().eq(Session::getUserSessionId, id).orderByAsc(Session::getId));
        //return sessionService.lambdaQuery().eq(Session::getUserSessionId, id).orderByAsc(Session::getId).list();
    }

    private void douBaoChatResult(SseEmitter emitter, List<ChatMessage> messages, Long userSessionId) throws Exception {
        Flowable<ChatCompletionChunk> result = ChatUtil.doubaoChat(messages);
        result.subscribe(
                message -> {
                    try {
                        String msg = message.getChoices().get(0).getMessage().getContent().toString();
                        String requestId = message.getId();
                        Session one = sessionService.lambdaQuery().eq(Session::getUserSessionId, userSessionId).eq(Session::getRequestId, requestId).one();
                        if (one != null) {
                            one.setContent(one.getContent() + msg);
                            sessionService.updateById(one);
                        } else {
                            Session save = new Session();
                            ChatMessageRole role = message.getChoices().get(0).getMessage().getRole();
                            save.setUserSessionId(userSessionId);
                            save.setRole(role.value());
                            save.setRequestId(requestId);
                            save.setContent(msg);
                            //save.setContent(AESUtil.encrypt(msg));//aes算法加密会话
                            sessionService.save(save);
                        }
                        emitter.send(msg);
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError, // 如果出错，终止流
                emitter::complete           // 完成整个流
        );
    }

    private void deepSeekChatResult(SseEmitter emitter, List<ChatMessage> messages, Long userSessionId) {
        Flux<String> stringFlux = DeepSeekUtil.streamChat(messages);
        stringFlux.subscribe(
                chunk -> {
                    try {
                        String role = "";
                        String msg = "";
                        // 解析并发送每个chunk到前端
//                        String msg = DeepSeekUtil.parseChunk(chunk);
                        JsonNode node = new ObjectMapper().readTree(chunk);
                        String requestId = node.get("id").asText();
                        if (node.get("choices").get(0).has("delta") && node.get("choices").get(0).get("delta").has("role")) {
                            role = node.get("choices").get(0).get("delta").get("role").asText();
                        }
                        if (node.get("choices").get(0).has("delta") && node.get("choices").get(0).get("delta").has("content")) {
                            msg = node.get("choices").get(0).get("delta").get("content").asText();
                        }
                        Session one = sessionService.lambdaQuery().eq(Session::getUserSessionId, userSessionId).eq(Session::getRequestId, requestId).one();
                        if (one != null) {
                            one.setContent(one.getContent() + msg);
                            sessionService.updateById(one);
                        } else {
                            Session save = new Session();
                            save.setUserSessionId(userSessionId);
                            save.setRole(role);
                            save.setRequestId(requestId);
                            save.setContent(msg);
                            //save.setContent(AESUtil.encrypt(msg));//aes算法加密会话
                            sessionService.save(save);
                        }
                        if (msg != null && !msg.isEmpty()) {
                            emitter.send(SseEmitter.event().data(msg));
                        }
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete
        );

        // 处理客户端断开连接
        emitter.onCompletion(() -> System.out.println("SSE connection completed"));
        emitter.onTimeout(() -> System.out.println("SSE connection timed out"));
        emitter.onError((ex) -> System.out.println("SSE error: " + ex.getMessage()));
    }

    private void chatGPTResult(SseEmitter emitter, List<ChatMessage> messages, Long userSessionId, String model) {
        Flux<String> stringFlux = ChatGPTUtil.streamChat(messages, model);
        stringFlux.subscribe(
                chunk -> {
                    try {
                        // 解析并发送每个chunk到前端
//                        String msg = ChatGPTUtil.parseChunk(chunk);
                        JsonNode node = new ObjectMapper().readTree(chunk);
                        String requestId = node.get("id").asText();
                        String role = "assistant";
                        String msg = node.get("choices").size() != 0 ? node.get("choices").get(0).get("delta").get("content").asText() : "";
                        if (!requestId.equals("") && !msg.equals("")){
                            Session one = sessionService.lambdaQuery().eq(Session::getUserSessionId, userSessionId).eq(Session::getRequestId, requestId).one();
                            if (one != null) {
                                one.setContent(one.getContent() + msg);
                                sessionService.updateById(one);
                            } else {
                                Session save = new Session();
                                save.setUserSessionId(userSessionId);
                                save.setRole(role);
                                save.setRequestId(requestId);
                                save.setContent(msg);
                                //save.setContent(AESUtil.encrypt(msg));//aes算法加密会话
                                sessionService.save(save);
                            }
                        }
                        if (msg != null && !msg.isEmpty()) {
                            emitter.send(SseEmitter.event().data(msg));
                        }
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete
        );

        // 处理客户端断开连接
        emitter.onCompletion(() -> System.out.println("SSE connection completed"));
        emitter.onTimeout(() -> System.out.println("SSE connection timed out"));
        emitter.onError((ex) -> System.out.println("SSE error: " + ex.getMessage()));
    }
}
